import os
import random
from PIL import Image


def parse_yolo_labels(label_path):
    boxes = []
    """
    yolo的txt格式
<class_id> <x_center> <y_center> <width> <height>   后四个都缩放到0～1
    """
    with open(label_path, 'r') as f:
        for line in f:
            parts = line.strip().split()
            if len(parts) == 5:
                cls, x, y, w, h = map(float, parts)
                # cls = 0
                boxes.append((cls, x, y, w, h))
    return boxes


def adjust_boxes(boxes, col, row, grid_size):
    adjusted = []
    for cls, x, y, w, h in boxes:
        new_x = (x + col) / grid_size
        new_y = (y + row) / grid_size
        new_w = w / grid_size
        new_h = h / grid_size
        adjusted.append((cls, new_x, new_y, new_w, new_h))
    return adjusted


def save_yolo_labels(boxes, save_path):
    with open(save_path, 'w') as f:
        for cls, x, y, w, h in boxes:
            f.write(f"{int(cls)} {x:.6f} {y:.6f} {w:.6f} {h:.6f}\n")


def create_mosaic_images(image_folder, label_folder, output_folder, grid_size, num_output_images, seed=42, start_count=0):
    # 设置固定随机种子
    random.seed(seed)

    os.makedirs(os.path.join(output_folder, "images"), exist_ok=True)
    os.makedirs(os.path.join(output_folder, "labels"), exist_ok=True)

    valid_exts = ('.jpg', '.jpeg', '.png', '.webp', 'bmp')
    image_files = [f for f in os.listdir(image_folder) if f.lower().endswith(valid_exts)]
    if not image_files:
        print("❌ 未找到图片文件。")
        return

    sample_img = Image.open(os.path.join(image_folder, image_files[0]))
    unit_w, unit_h = sample_img.size

    total_needed = grid_size * grid_size * num_output_images
    replicated_list = (image_files * ((total_needed // len(image_files)) + 1))[:total_needed]
    random.shuffle(replicated_list)

    for idx in range(num_output_images):
        mosaic_img = Image.new('RGB', (unit_w * grid_size, unit_h * grid_size))
        mosaic_labels = []

        selected_files = replicated_list[idx * grid_size * grid_size:(idx + 1) * grid_size * grid_size]

        for i, img_file in enumerate(selected_files):
            col = i % grid_size
            row = i // grid_size
            img_path = os.path.join(image_folder, img_file)
            label_path = os.path.join(label_folder, os.path.splitext(img_file)[0] + '.txt')

            assert os.path.exists(label_path), f"❌ 标签文件缺失: {label_path}"

            img = Image.open(img_path).convert("RGB")
            mosaic_img.paste(img, (col * unit_w, row * unit_h))

            boxes = parse_yolo_labels(label_path)
            adj_boxes = adjust_boxes(boxes, col, row, grid_size)
            mosaic_labels.extend(adj_boxes)

        mosaic_img = mosaic_img.resize((unit_w, unit_h))
        image_name = f"test_{start_count+idx}.png"
        label_name = f"test_{start_count+idx}.txt"
        mosaic_img.save(os.path.join(output_folder, "images", image_name), format="PNG")
        save_yolo_labels(mosaic_labels, os.path.join(output_folder, "labels", label_name))

    print(f"✅ 成功生成 {num_output_images} 张 {grid_size}x{grid_size} 拼图，保存在 {output_folder}。")


if __name__ == "__main__":

    grid_sizes = [2, 3, 4]
    start_num = 0  # 下次遍历生成图片的起始索引
    for grid_size in grid_sizes:
        root_folder = r'F:\project\526\fake_image\test\test_split'  # 各种分辨率的图片的上层地址
        subfolders = [f for f in os.listdir(root_folder) if os.path.isdir(os.path.join(root_folder, f))]
        for subfolder in subfolders:
            image_folder = os.path.join(root_folder, subfolder)
            # 获取该文件夹下图片的数量
            image_files = [f for f in os.listdir(image_folder) if f.lower().endswith(('jpg', 'jpeg', 'png', 'webp', 'bmp'))]
            num_images = len(image_files)
            num_concate_images = 0
            if num_images < grid_size * grid_size:
                num_concate_images = 1
            else:
                num_concate_images = (num_images // (grid_size * grid_size)) * 3
            print(num_concate_images)
            # if num_concate_images > 150:
            #     num_concate_images = 150
            create_mosaic_images(
                image_folder=image_folder,  # 输入图像目录
                label_folder=r"F:\project\526\fake_image\test\test_label",  # 输入标签目录
                output_folder=r"F:\project\526\fake_image\test\test_merge",  # 输出路径
                grid_size=grid_size,  # 3x3 拼接
                num_output_images=num_concate_images,  # 生成 10 张
                seed=2025,  # 控制打乱顺序
                start_count=start_num  # 控制图片下标索引从几开始
            )
            start_num += num_concate_images





